#include "cudaCom.h"
#include <glm/glm.hpp>
#include "lbcuda.h"
#include <cstring>



static float cpuWeight[9] = {
		4.0f/9.0f,

		1.0f/9.0f,
		1.0f/9.0f,
		1.0f/9.0f,
		1.0f/9.0f,

		1.0f/36.0f,
		1.0f/36.0f,
		1.0f/36.0f,
		1.0f/36.0f
	};


static int cpuE[18] = {
		0,0,1, 0,-1,1, 1,-1,-1, //x
		0,1,0,-1, 0,1,-1,-1, 1  //y
	};

float cpuVisc = 0.8f;

__constant__ int numCellsU;
__constant__ int numCellsV;
__constant__ int e[18]; // 9 x's followd by 9 y's
__constant__ float weight[9];
__constant__ float VISC;

__device__ void GetCellUV(int cndx, int & outU, int & outV){
	outU = cndx % numCellsU;
	outV = cndx / numCellsU;
}

__device__ int GetCellNdx(int cellU, int cellV){
	cellU = cellU >= numCellsU  ? numCellsU-1 : cellU;
	cellU = cellU < 0 ? 0 : cellU;
	cellV = cellV >= numCellsV  ? numCellsV-1 : cellV;
	cellV = cellV < 0 ? 0 : cellV;

	return (cellV*numCellsU + cellU);
}


__device__ float CalcDensity(int cndx, float * lbCells ){
	float sum = 0.0f;
	for (int i = cndx; i < cndx+9; ++i){
		sum+= lbCells[i];
	}
	return sum;
}
__device__ void CalcVel(int cndx, float * lbField, float & outVx, float & outVy) {
	outVx = outVy = 0.0f;

	for (int i = 0; i < 9; ++i){
		float fi = lbField[cndx+i];
		outVx += e[i]*fi;
		outVy += e[i+9]*fi;
	}

	return;
}

__device__ float CalcFeq(int i, int cndx, float dens, float velx, float vely){
	float edotu = velx*e[i] + vely*e[i+9];
	float vdot = velx*velx + vely*vely;

	return
		weight[i]*
		(dens +
		3 * edotu - 
		(3.0f/2.0f)*(vdot) +
		(9.0f/2.0f)*(pow(edotu, 2.0f) ) );
}


__global__ void kStreamCells(float * lbCellsSrc, float * lbCellsDst) {
	unsigned int cndx = blockIdx.x*blockDim.x + threadIdx.x;
	int NUMCELLS = numCellsU*numCellsV;


	if (cndx >= NUMCELLS){ return; }



	// From cndx, get cellu, cellv
	int cellU, cellV;
	GetCellUV(cndx, cellU, cellV);
	int offcellU, offcellV;

	cndx *= 9;

	//Stream
	for (int i = 0; i < 9; ++i){
		int offset[2] = {e[i], e[i+9]};

		// Minus to go backwards
		offcellU = cellU - offset[0];
		offcellV = cellV - offset[1];

		// periodic
		//if (PERIODIC){
			if (offcellU >= numCellsU){ offcellU = 0; }
			if (offcellU < 0){ offcellU = numCellsU-1; }
			if (offcellV >= numCellsV) { offcellV = 0; }
			if (offcellV < 0){ offcellV = numCellsV-1; }
	//	} 

	
		int offndx = GetCellNdx(offcellU, offcellV)*9;
		//if (PERIODIC){
		
			lbCellsDst[cndx+i]  = lbCellsSrc[offndx+i];
		//	lbCellsDst[0] = lbCellsSrc[0];
		
		//} else {
		//	// Dealing with boundary cell in fi direction
		//	
		//	if (lbFieldObs[offndx] && i != 0){
		//		lbfield[(currNdx+1)%2][cndx][i] = lbfield[currNdx][offndx][ ((i+1)%4)+ (i < 5 ? 1 : 5) ] ;
		//	} else {
		//		lbfield[(currNdx+1) % 2][cndx][i]  = lbfield[currNdx][offndx][i];				
		//	}
		//}
	}
}

__global__ void kUpdateParts(float * partsArry, float * grid, int numParts ){
	unsigned int pndx = blockIdx.x*blockDim.x + threadIdx.x;
	if (pndx >= numParts){
		return;
	}

	float * part = &partsArry[pndx*2];

	// Get grid cell and calculate velocity
	int cndx = GetCellNdx(floor(part[0]), floor(part[1])) * 9;

	// Get the velocity.
	float vel[2];
	CalcVel(cndx, grid, vel[0], vel[1]);

	part[0] += vel[0];
	part[1] += vel[1];

	return;
}

__global__ void kCollideCells( float *lbCellSrc, float * lbCellDst){
	unsigned int cndx = blockIdx.x*blockDim.x + threadIdx.x;
	int NUMCELLS = numCellsU*numCellsV;

	if (cndx >= NUMCELLS){ return; }
	// From cndx, get cellu, cellv
	int cellU, cellV;
	GetCellUV(cndx, cellU, cellV);
	cndx *= 9;

	float dens = CalcDensity(cndx, lbCellSrc);
	float vel[2];
	CalcVel(cndx, lbCellSrc, vel[0], vel[1]);

	for (int i = 0; i < 9; ++i){
		float feq = CalcFeq(i, cndx, dens, vel[0], vel[1]);

		lbCellDst[cndx+i] = 
			(1.0f-VISC) * lbCellDst[cndx+i] + 
			VISC * feq;
	}
}

//// Host


CudaLB::CudaLB(int numCellsU, int numCellsV)
{
	isAllocated = false;
	mNumCells[0] = numCellsU;
	mNumCells[1] = numCellsV;
	mNumParts = 0;
	currNdx = 0;

	AllocateMemory();
	
	cudaDeviceSynchronize();

	isAllocated = true;
}

void CudaLB::AllocateConstants() {

	
	

	DBG_CUDA_CALL(
		cudaMemcpyToSymbol("numCellsU", &mNumCells[0] , sizeof(int)),
		"Error copying numCellsU to constant memory");

	DBG_CUDA_CALL(
		cudaMemcpyToSymbol("numCellsV", &mNumCells[1] , sizeof(int)),
		"Error copying numCellsV to constant memory");

	DBG_CUDA_CALL(
		cudaMemcpyToSymbol("weight", &cpuWeight[0] , sizeof(float)*9),
		"Error copying weights to constant memory");
	
	DBG_CUDA_CALL(
		cudaMemcpyToSymbol("e", &cpuE[0] , sizeof(int)*18),
		"Error copying e to constant memory");

	DBG_CUDA_CALL(
		cudaMemcpyToSymbol("VISC", &(cpuVisc), sizeof(float)),
		"Error copying visc to constant memory");

}

void CudaLB::AllocateMemory(){

	// Allocate enough memory for the LB grid.
	size_t sizeGrid = 
		sizeof(float)*mNumCells[0]*mNumCells[1]*9;
	
	for (int i = 0; i < 2; ++i){
		DBG_CUDA_CALL(
			cudaMalloc((void**)(&(lbGrid[i])), sizeGrid),
			"Error allocating lbGrid");
		DBG_CUDA_CALL(
			cudaMemset((void*)lbGrid[i], 0, sizeGrid),
			"Error initting memory on grid" );
	}

	// Allocates the constants
	AllocateConstants();

}

void CudaLB::GetGrid( float * outMem ){

	cudaDeviceSynchronize();

	DBG_CUDA_CALL(
		cudaMemcpy(outMem, lbGrid[currNdx], sizeof(float) * mNumCells[0]*mNumCells[1]*9, cudaMemcpyDeviceToHost),
		"Trying to transfer lbgrid from device to host");

	return;

}

void CudaLB::SetGrid( float * inmem ){

	cudaDeviceSynchronize();

	DBG_CUDA_CALL(
		cudaMemcpy(lbGrid[currNdx], inmem, sizeof(float) * mNumCells[0]*mNumCells[1]*9, cudaMemcpyHostToDevice),
		"Trying to transfer lbgrid from device to host");

	return;

}

void CudaLB::Update(){


	cudaDeviceSynchronize();

	// Calls the stream and collider kernels.

	// Num of particles
	int numThreads = 512;
	int numBlocks = ((mNumCells[0]*mNumCells[1])/numThreads) + 1;

	//printf("NumThreads : %d, numBlocks: %d\n", numThreads, numBlocks);

	kStreamCells<<<numBlocks, numThreads>>> (lbGrid[currNdx], lbGrid[(currNdx+1)%2]);
	cudaDeviceSynchronize();
	kCollideCells<<<numBlocks, numThreads>>> (lbGrid[currNdx], lbGrid[(currNdx+1)%2]);
	cudaDeviceSynchronize();
	


	//printf("Checking error...");
//	DBG_CUDA_CHECK_ERRORS
	
	// Sets the current ndx	
	currNdx = (currNdx+1)%2;

}

void CudaLB::UpdateParts( ){
	if (mNumParts != 0){
		float * vboPtr;

		MapCudaVBO(mPartsVBO, &vboPtr);
	
		int numThreads = 256;
		int numBlocks = ((mNumParts)/numThreads) + 1;

		kUpdateParts<<<numBlocks, numThreads>>> (vboPtr, lbGrid[currNdx], mNumParts);

		cudaDeviceSynchronize();


		UnmapCudaVBO(mPartsVBO);
	}

}

void CudaLB::SetNumParts( int num, int dim, float * parts){
	// Allocates the parts.
	// Delete current VBO
	DeleteCUDAGLVBO( mPartsVBO );

	// Generate a new one
	mPartsVBO = GenerateCUDAGLVbo( num, dim, parts );

	// Register cuda vbo
	RegisterCudaVBO(mPartsVBO);

	mNumParts = num;
}

int CudaLB::GetNumParts() {
	return mNumParts;
}

GLuint CudaLB::GetPartsVBO() {
	return mPartsVBO;
}

// This code can be used for opengl interop.

// CUDA OpenGL Interop stuff

// Initialize the cudagl device.  This is nessesary if doing any kind of 
// OpenGL-cuda interop.
void InitCudaGLDevice() {
	DBG_CUDA_CALL(cudaGLSetGLDevice(0),
		"Error setting the GL device to 0");
}

GLuint GenerateCUDAGLVbo( int numVerts, int numDims, float * data ){
	GLuint outVBO;

	glGenBuffers(1, &outVBO);
	glBindBuffer(GL_ARRAY_BUFFER, outVBO);

	// Initialize vbo
	unsigned int size = sizeof(float)*numDims*numVerts;
	glBufferData(GL_ARRAY_BUFFER, size, data, GL_DYNAMIC_DRAW);

	// Unbind VBO
	glBindBuffer(GL_ARRAY_BUFFER, 0);

	return outVBO;
}

void DeleteCUDAGLVBO( GLuint vbo ){
	glDeleteBuffers(1, &vbo);

	return;
}

void RegisterCudaVBO( GLuint vbo ){

	// Register buffer object with CUDA
	DBG_CUDA_CALL(
		cudaGLRegisterBufferObject(vbo),
		"Error trying to register buffer object");
}

void MapCudaVBO ( GLuint vbo, float ** mappedPtr ){
	DBG_CUDA_CALL( 
		cudaGLMapBufferObject( (void**)mappedPtr, vbo),
		"Error trying to map buffer object with CUDA.");
}

void RenderVBO( GLuint vbo, int numDims, int numPoints, float * color ){
	glBindBuffer(GL_ARRAY_BUFFER, vbo);
	glVertexPointer(numDims, GL_FLOAT, 0, 0);

	glEnableClientState(GL_VERTEX_ARRAY);
	glColor3fv(color);
	glDrawArrays(GL_POINTS, 0, numDims*numPoints);
	glDisableClientState(GL_VERTEX_ARRAY);

	glBindBuffer(GL_ARRAY_BUFFER, 0);
}

void UnmapCudaVBO( GLuint vbo ){
	DBG_CUDA_CALL(
		cudaGLUnmapBufferObject( vbo ),
		"Error unmapping VBO" );
}


//// Create an opengl texture
//void CreateEmptyTexture(int resU, int resV, GLuint & tex) {
//	glGenTextures(1, &tex);
//	glBindTexture(GL_TEXTURE_2D, tex);
//
//	// Set Params
//	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE);
//	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE);
//	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
//	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
//
//	//Create the texture data
//	glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA,
// 
//}